package com.jboltai.capability.embedding.ali;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.TypeReference;
import com.jboltai.capability.embedding.EmbeddingProcessor;
import com.jboltai.config.AIParamKey;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.EventAttrKey;
import com.jboltai.event.EventErrorType;
import com.jboltai.event.EventState;
import com.jboltai.event.embedding.EmbeddingEvent;
import com.jboltai.resource.ai.AIResource;
import com.jboltai.resource.embedding.EmbeddingResource;
import com.jboltai.resource.embedding.EmbeddingResourceCenter;
import com.jboltai.util.http.HttpUtil;
import com.jboltai.util.http.RequestInstance;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class TongyiEmbeddingImpl implements EmbeddingProcessor {

    public static final TongyiEmbeddingImpl INSTANCE = new TongyiEmbeddingImpl();
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    public static final String API_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings";

    private static final TypeReference TYPE = new TypeReference<List<Double>>() {};
    @Override
    public void process(EmbeddingEvent event, EmbeddingResource resource) {
        LOGGER.debug("开始 调用通义Embedding大模型 处理Embedding事件：{}", event.getId());
        //资源指定了apiurl就用资源的，否则用默认的
        String apiUrl = StrUtil.isNotBlank(resource.getApiUrl()) ? resource.getApiUrl() : API_URL;
        RequestInstance request = HttpUtil.post(resource.getClient(), apiUrl, event);
        request.setAuthorization(resource.getApiKey());
        request.addJsonParam(AIParamKey.MODEL, resource.getDefaultModelName());
        request.addJsonParam("input", event.getText());

        event.setAttr(EventAttrKey.REQUEST_INSTANCE,request);

        request.asyncSend(response -> {
            try {
                String content = response.body().string();
                LOGGER.debug("Embedding[{}]请求结束:{}", event.getId(), content);
                if (response.code() != 200) {
                    //有异常情况
                    handleFail(event, resource, response.code(), content, null);
                } else {
                    JSONObject data = JSON.parseObject(content);
                    JSONArray list = data.getJSONArray("data");
                    if (list == null) {
                        event.success(Collections.emptyList());
                    } else {
                        List<List<Double>> result = new ArrayList<>();
                        for (int i = 0; i < list.size(); i++) {
                            List<Double> item = list.getJSONObject(i).getObject("embedding", TYPE);
                            result.add(item);
                        }
                        event.setResult(result);
                        event.success(result);
                    }
                    event.setState(EventState.SUCCESS);
                }
            } catch (Exception e) {
                LOGGER.error("Embedding[{}]处理响应结果发生异常:{}", event.getId(), e);
                event.fail(EventErrorType.PROGRAM_EXCEPTION, e);
            } finally {
                event.complete();
            }
        }, e -> {
            LOGGER.error("Embedding[{}]请求失败:{}", event.getId(), e);
            event.fail(EventErrorType.NET_ERROR, e);
        });
    }


    /**
     * 失败时的处理
     * @param event
     * @param resource
     * @param code
     * @param responseContent
     * @param exception
     */
    private static void handleFail(EmbeddingEvent event, EmbeddingResource resource, int code, String responseContent, Throwable exception) {
        JSONObject error  = null;;
        LOGGER.error("EmbeddingEvent[{}]请求失败:{}_{},异常：{}", event.getId(), code, responseContent, exception);
        if (exception == null) {
            //ai服务异常
            switch (code) {
                case 400:{
                    error = JSON.parseObject(responseContent);
                    event.fail(EventErrorType.REQUEST_PARAM_ERROR, error.getJSONObject("error").getString("message"));
                    break;
                }

                case 401:
                    event.fail(EventErrorType.API_KEY_INVALID, responseContent);
                    break;
                case 429:{
                    error = JSON.parseObject(responseContent);
                    if ("rate_limit".equals(error.getJSONObject("error").getString("code"))) {
                        EmbeddingResourceCenter.moveResourceToInvalid(resource);
                        event.fail(EventErrorType.EXCEED_REQUEST_LIMIT, responseContent);
                    } else {
                        event.fail(EventErrorType.ACCOUNT_BALANCE_INSUFFICIENT, responseContent);
                    }
                    break;
                }
                case 500:
                case 503:
                    event.fail(EventErrorType.SERVICE_SERVER_ERROR);
                    break;
                default:
                    //未知错误
                    event.fail(EventErrorType.UNKNOWN,responseContent);
                    break;
            }
        } else {
            //程序异常
            event.fail(EventErrorType.NET_ERROR, exception);
        }
    }

    @Override
    public void handleCancel(EmbeddingEvent event, AIResource resource) {
        LOGGER.warn("{}事件取消执行", event.getId());
        RequestInstance requestInstance = event.getAttr(EventAttrKey.REQUEST_INSTANCE);
        requestInstance.cancal();
    }
}
